from skorch.core import *
from skorch.layers import *
from skorch.vision import *
from skorch.callbacks.hooks import *
from skorch.model import create_cnn


def unet(n_classes: int=2, feature_scale:int=1, imsize:Sizes = (256,256), arch:Callable = models.resnet18, ce_encoder=False,
         pretrained:bool=True, cut: Union[int, Callable] = None, debug=False):
    model = models.OriginalUNet(n_classes=n_classes, feature_scale=feature_scale, ce_encoder=ce_encoder,debug=debug)
    return model


def classify_cnn(n_classes: int=2, softmax: bool=False, feature_scale:int=1, imsize:Sizes=(256,256)):
    """
    classification module
    :param n_classes: classes number
    :param softmax: add softmax layer
    :param feature_scale: scale value
    :param imsize: image size
    :return:
    """
    model = create_cnn(models.resnet18, nc=n_classes, softmax=softmax, pretrained=False)#, y_range=[0,1])
    return model



def main():
    #model = unet(feature_scale=1, n_classes=3, debug=True)
    model = create_cnn(models.resnet18, nc=4)
    print(model_sizes(model, size=(512, 512)))

    x = torch.ones((1, 3, 512, 512))
    y = model.forward(x)


    #summary(net, (3,512,512), batch_size=1)


if __name__ == '__main__':
    main()